#里面的文件夹名字直接hardcode进去了，大部分目标检测算法都是这样划分的
#在经过split_train_val_yolo.py处理后，我们划分出了训练集、测试集
#文件夹排版
# ├── Annotations
    #标注的xml文件
# ├── images
    #图片
# └── ImagesSets
#     └── Main
         #train.txt、trainval.txt、val.txt
#经过该脚本后  需要修改标签(classes)和root-dir
# ├── Annotations
    #标注的xml文件
# ├── dataSet_path
    #划分出的训练集和验证集
# ├── images
    #图片
# ├── ImagesSets
# │   └── Main
         #train.txt、trainval.txt、val.txt
# └── labels
    #yolo格式的图片
import xml.etree.ElementTree as ET
import argparse
from pathlib import Path

def Parse_Arguments():
    parser = argparse.ArgumentParser(description="")
    parser.add_argument('--root-dir', type=str, default="/home/luoluoluo/data/dataset/elevator/train_detect")
    return parser.parse_args()

sets = ['train', 'val']
args = Parse_Arguments()
classes = []
with open(args.root_dir + "/classes.txt", 'r') as f:
    classes = f.read().strip().splitlines()


#矩形框 (由左上角和右下角的坐标表示) 改为 (中心点坐标和宽度、高度)
def convert(size, box):
    dw = 1. / (size[0]+0.01)
    dh = 1. / (size[1]+0.01)
    # x = (box[0] + box[1]) / 2.0 - 1
    # y = (box[2] + box[3]) / 2.0 - 1
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h


def convert_annotation(image_id):
    in_file = open(f"{args.root_dir}/Annotations/{image_id}.xml", encoding='UTF-8')
    out_file = open(f"{args.root_dir}/labels/{image_id}.txt", 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        if obj.find('difficult'):
            difficult = obj.find('difficult').text
        else:
            difficult = 0
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        b1, b2, b3, b4 = b
        # 标注越界修正
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = (b1, b2, b3, b4)
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


def main():
    for image_set in sets:
        if not Path(f"{args.root_dir}/labels/").exists():
            Path(f"{args.root_dir}/labels/").mkdir()
        image_ids = open(f"{args.root_dir}/ImagesSets/Main/{image_set}.txt").read().strip().split()
        if not Path(f"{args.root_dir}/dataSet_path/").exists():
            Path(f"{args.root_dir}/dataSet_path/").mkdir()

        list_file = open(f"{args.root_dir}/dataSet_path/{image_set}.txt", 'w')
        for image_id in image_ids:
            list_file.write(f"{args.root_dir}/images/{image_id}.jpg\n")
            convert_annotation(image_id)
        list_file.close()
        print("xml-to-yolo done!")
if __name__ == "__main__":
    main()
